// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef _CahnHilliard_EquationSet_impl_hpp_
#define _CahnHilliard_EquationSet_impl_hpp_

#include "Teuchos_ParameterList.hpp"
#include "Teuchos_StandardParameterEntryValidators.hpp"
#include "Teuchos_RCP.hpp"
#include "Teuchos_Assert.hpp"
#include "Phalanx_DataLayout_MDALayout.hpp"
#include "Phalanx_FieldManager.hpp"

#include "Panzer_IntegrationRule.hpp"
#include "Panzer_BasisIRLayout.hpp"

// include evaluators here
#include "Panzer_Integrator_BasisTimesScalar.hpp"
#include "Panzer_Integrator_TransientBasisTimesScalar.hpp"
#include "Panzer_Integrator_BasisTimesVector.hpp"
#include "Panzer_Integrator_BasisTimesTensorTimesVector.hpp"
#include "Panzer_Product.hpp"

#include "CahnHillChemTerm.hpp"


// ***********************************************************************
template <typename EvalT>
SiYuan::EquationSet_CahnHilliard<EvalT>::
EquationSet_CahnHilliard(const Teuchos::RCP<Teuchos::ParameterList>& params,
    const int& default_integration_order,
    const panzer::CellData& cell_data,
    const Teuchos::RCP<panzer::GlobalData>& global_data,
    const bool build_transient_support) :
    EquationSet<EvalT> (params,default_integration_order,cell_data,global_data,build_transient_support ) 
{
  std::string model_id = params->get<std::string>("Model ID");
  m_dof_names[0] = "Concentration"; // The concentration difference variable
  m_dof_names[1] = "Diffusion Potential";     // The chemical potential 
  std::string basis_type = "HGrad";
  int basis_order = params->get<int>("Basis Order");
  int integration_order = params->get<int>("Integration Order");
  beta_ = params->get<double>("Beta");

  // ********************
  // Setup DOFs and closure models
  // ********************
  {
    this->addDOF(m_dof_names[0],basis_type,basis_order,integration_order);
    this->addDOFGrad(m_dof_names[0]);
    this->addDOFTimeDerivative(m_dof_names[0]);

    this->addDOF(m_dof_names[1],basis_type,basis_order,integration_order);
    this->addDOFGrad(m_dof_names[1]);
    this->addDOFTimeDerivative(m_dof_names[1]);   // Is it needed?
  }

  this->addClosureModel(model_id);
  this->setupDOFs();

  this->addC0("Diffusion Mobility");
  this->addC3("Gradient Energy Coefficient");
}

// ***********************************************************************
template <typename EvalT>
void SiYuan::EquationSet_CahnHilliard<EvalT>::
buildAndRegisterEquationSetEvaluators(PHX::FieldManager<panzer::Traits>& fm,
    const panzer::FieldLibrary& /* fl */,
    const Teuchos::ParameterList& ) const
{
  using panzer::EvaluatorStyle;

  Teuchos::RCP<panzer::IntegrationRule> ir = this->getIntRuleForDOF(m_dof_names[0]); 
  Teuchos::RCP<panzer::BasisIRLayout> basis = this->getBasisIRLayoutForDOF(m_dof_names[0]);

  // Residual: Concentration
  {
    std::string resName("RESIDUAL_Concentration");
    std::string valName( "DXDT_" + m_dof_names[0] );
    std::string gradName("GRAD_" + m_dof_names[1]);
    double multiplier(1);
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> top =Teuchos:: rcp(new
      SiYuan::Residual_Concentration<EvalT, panzer::Traits>(EvaluatorStyle::EVALUATES,
        resName, valName, gradName, *basis, *ir, this->getC0() ) );
    this->template registerEvaluator<EvalT>(fm, top);
  }
/*  {
    std::string resName("RESIDUAL_" + m_dof_names[0]);
    std::string valName("DXDT_" + m_dof_names[0]);
    double multiplier(1);
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> top =Teuchos:: rcp(new
      panzer::Integrator_BasisTimesScalar<EvalT, panzer::Traits>(EvaluatorStyle::EVALUATES,
        resName, valName, *basis, *ir, multiplier) );
    this->template registerEvaluator<EvalT>(fm, top);
  }

  // Diffusion Operator
  {
    Teuchos::ParameterList p("Diffusion Residual");
    p.set("Residual Name", "RESIDUAL_"+m_dof_names[0]);
    p.set("Flux Name", "GRAD_"+m_dof_names[1]);
    p.set("Basis", basis);
    p.set("IR", ir);
    Teuchos::RCP<const std::vector<std::string> > vec = 
      Teuchos::rcp(new std::vector<std::string>(this->getC0()));
    p.set("Field Multipliers", vec);
    p.set("Multiplier", 1.0);
    
    Teuchos::RCP< PHX::Evaluator<panzer::Traits> > op = 
      Teuchos::rcp(new panzer::Integrator_GradBasisDotVector<EvalT,panzer::Traits>(p));

    this->template registerEvaluator<EvalT>(fm, op);
  }
*/
  // Chemistry Term
  {
    Teuchos::ParameterList p("Chemistry Term");
    p.set("Beta", beta_);
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> top =Teuchos:: rcp(new
      SiYuan::CahnHillChemTerm<EvalT, panzer::Traits>(p, *ir ) );
    this->template registerEvaluator<EvalT>(fm, top);
  }

  Teuchos::RCP<panzer::IntegrationRule> ir1 = this->getIntRuleForDOF(m_dof_names[1]); 
  Teuchos::RCP<panzer::BasisIRLayout> basis1 = this->getBasisIRLayoutForDOF(m_dof_names[1]);

  // Residual: Chemistry Potential
  {
    std::string resName("RESIDUAL_Diffusion Potential");
    std::string valName( "FreeEnergy" );
    std::string gradName("GRAD_" + m_dof_names[0]);
    double multiplier(1);
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> top =Teuchos:: rcp(new
      SiYuan::Residual_ChemistryPotential<EvalT, panzer::Traits>(EvaluatorStyle::EVALUATES,
        resName, valName, gradName, *basis1, *ir1, this->getC3() ) );
    this->template registerEvaluator<EvalT>(fm, top);
  }

  // Use a sum operator to form the overall residual for the equation
 /* {
    std::vector<std::string> sum_names;
    
    // these are the names of the residual values to sum together
    sum_names.push_back("RESIDUAL_Concentration");
    sum_names.push_back("RESIDUAL_Diffusion");

    this->buildAndRegisterResidualSummationEvaluator(fm,"SUMCH",sum_names);
  }*/

}

// ***********************************************************************

// ***********************************************************************
/////////////////////////////////////////////////////////////////////////////
//
//  Main Constructor: Residual_ChemistryPotential (Rho residual)
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
SiYuan::Residual_ChemistryPotential<EvalT, Traits>::
Residual_ChemistryPotential(
    const panzer::EvaluatorStyle&   evalStyle,
    const std::string&              resName,
    const std::string&              valName,
    const std::string&              gradName,
    const panzer::BasisIRLayout&    basis,
    const panzer::IntegrationRule&  ir ,
    const std::vector<std::string>& c3 )
    :
    evalStyle_(evalStyle), basisName_(basis.name())
{
    cGrad_ = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP, panzer::Dim>(gradName, ir.dl_vector);
    this->addDependentField(cGrad_);
    m_freeEnergy = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP>(valName, ir.dl_scalar);
    this->addDependentField(m_freeEnergy);

    // Add the Material properties
    std::string mname = "Gradient Energy Coefficient";
    gamma_ = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP>(mname, ir.dl_scalar);
    this->addDependentField(gamma_);

    // Create the field that we're either contributing to or evaluating
    // (storing).
    field_ = PHX::MDField<ScalarT, panzer::Cell, panzer::BASIS>(resName, basis.functional);
    if (evalStyle == panzer::EvaluatorStyle::CONTRIBUTES)
      this->addContributedField(field_);
    else // if (evalStyle == EvaluatorStyle::EVALUATES)
      this->addEvaluatedField(field_);

    // Set the name of this object.
    std::string n("Integrator_ChemistryPotential (");
    if (evalStyle == panzer::EvaluatorStyle::CONTRIBUTES)
      n += "CONTRIBUTES";
    else // if (evalStyle == EvaluatorStyle::EVALUATES)
      n += "EVALUATES";
    n += "):  " + field_.fieldTag().name();
    this->setName(n);
} // end of Main Constructor

/////////////////////////////////////////////////////////////////////////////
//
//  postRegistrationSetup()
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
void
SiYuan::Residual_ChemistryPotential<EvalT, Traits>:: postRegistrationSetup(
    typename Traits::SetupData sd, PHX::FieldManager<Traits>& fm)
{
    basisIndex_ = panzer::getBasisIndex(basisName_, (*sd.worksets_)[0], this->wda);
} // end of postRegistrationSetup()

/////////////////////////////////////////////////////////////////////////////
//
//  evaluateFields()
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
void
SiYuan::Residual_ChemistryPotential<EvalT, Traits>::
evaluateFields( typename Traits::EvalData workset )
{
    using Kokkos::parallel_for;
    using Kokkos::RangePolicy;
    using PHX::Device;

    typedef Intrepid2::FunctionSpaceTools<PHX::exec_space> FST;

//    if( workset.beta==0.0 ) return;
//std::cout << "RES:" << PHX::print<EvalT>()  << std::endl;
    // Grab the basis information.
    const panzer::BasisValues2<double>& bv = *this->wda(workset).bases[basisIndex_];

    PHX::MDField<double, panzer::Cell, panzer::BASIS, panzer::IP, panzer::Dim>
          weightedGradBasis = bv.weighted_grad_basis;
    PHX::MDField<double, panzer::Cell, panzer::BASIS, panzer::IP>
          weightedBasis = bv.weighted_basis_scalar;
 
    int numDims = weightedGradBasis.extent(3);
    int numBases= weightedGradBasis.extent(1);
    int numQP = weightedGradBasis.extent(2);
    for( int cell=0; cell<workset.num_cells; cell++ )
    {
        for (int basis(0); basis < numBases; ++basis)
        {
        //  if (evalStyle_ == panzer::EvaluatorStyle::EVALUATES)
          field_(cell, basis) = 0.0;
          for (int qp(0); qp < numQP; ++qp) {
            for (int dim(0); dim<numDims; ++dim ) {
              field_(cell, basis) += gamma_(cell, qp)* cGrad_(cell, qp, dim) *
                weightedGradBasis(cell, basis, qp, dim);
            }
            field_(cell, basis) += m_freeEnergy(cell, qp)* weightedBasis(cell, basis, qp);
        //    std::cout << qp << "," << m_freeEnergy(cell, qp) << ", " << weightedBasis(cell, basis, qp) << std::endl;
          }
        //  std::cout << cell << "," << basis << ", " << field_(cell, basis) << std::endl;
        } // end loop over the basis functions
    }

} // end of evaluateFields()


// ***********************************************************************

// ***********************************************************************
/////////////////////////////////////////////////////////////////////////////
//
//  Main Constructor: Residual_Concentration (W residual)
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
SiYuan::Residual_Concentration<EvalT, Traits>::
Residual_Concentration(
    const panzer::EvaluatorStyle&   evalStyle,
    const std::string&              resName,
    const std::string&              valName,
    const std::string&              gradName,
    const panzer::BasisIRLayout&    basis,
    const panzer::IntegrationRule&  ir ,
    const std::vector<std::string>& c3 )
    :
    evalStyle_(evalStyle), basisName_(basis.name())
{
  //  mu_ = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP>(valName, ir.dl_scalar);
  //  this->addDependentField(mu_);
    muGrad_ = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP, panzer::Dim>(gradName, ir.dl_vector);
    this->addDependentField(muGrad_);
    cdot_ = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP>(valName, ir.dl_scalar);
    this->addDependentField(cdot_);

    // Add the Material properties
    std::string mname = "Diffusion Mobility";
    M_diff_ = PHX::MDField<const ScalarT, panzer::Cell, panzer::IP>(mname, ir.dl_scalar);
    this->addDependentField(M_diff_);

    // Create the field that we're either contributing to or evaluating
    // (storing).
    field_ = PHX::MDField<ScalarT, panzer::Cell, panzer::BASIS>(resName, basis.functional);
    if (evalStyle == panzer::EvaluatorStyle::CONTRIBUTES)
      this->addContributedField(field_);
    else // if (evalStyle == EvaluatorStyle::EVALUATES)
      this->addEvaluatedField(field_);

    // Set the name of this object.
    std::string n("Integrator_Concentration (");
    if (evalStyle == panzer::EvaluatorStyle::CONTRIBUTES)
      n += "CONTRIBUTES";
    else // if (evalStyle == EvaluatorStyle::EVALUATES)
      n += "EVALUATES";
    n += "):  " + field_.fieldTag().name();
    this->setName(n);
} // end of Main Constructor

/////////////////////////////////////////////////////////////////////////////
//
//  postRegistrationSetup()
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
void
SiYuan::Residual_Concentration<EvalT, Traits>:: postRegistrationSetup(
    typename Traits::SetupData sd, PHX::FieldManager<Traits>& fm)
{
    basisIndex_ = panzer::getBasisIndex(basisName_, (*sd.worksets_)[0], this->wda);
} // end of postRegistrationSetup()

/////////////////////////////////////////////////////////////////////////////
//
//  evaluateFields()
//
/////////////////////////////////////////////////////////////////////////////
template<typename EvalT, typename Traits>
void
SiYuan::Residual_Concentration<EvalT, Traits>::
evaluateFields( typename Traits::EvalData workset )
{
    using Kokkos::parallel_for;
    using Kokkos::RangePolicy;
    using PHX::Device;

    typedef Intrepid2::FunctionSpaceTools<PHX::exec_space> FST;

//    if( workset.beta==0.0 ) return;
//std::cout << "RES:" << PHX::print<EvalT>()  << std::endl;
    // Grab the basis information.
    const panzer::BasisValues2<double>& bv = *this->wda(workset).bases[basisIndex_];

    PHX::MDField<double, panzer::Cell, panzer::BASIS, panzer::IP, panzer::Dim>
          weightedGradBasis = bv.weighted_grad_basis;
    PHX::MDField<double, panzer::Cell, panzer::BASIS, panzer::IP>
          weightedBasis = bv.weighted_basis_scalar;
 
    int numDims = weightedGradBasis.extent(3);
    int numBases= weightedGradBasis.extent(1);
    int numQP = weightedGradBasis.extent(2);
    for( int cell=0; cell<workset.num_cells; cell++ )
    {
        for (int basis(0); basis < numBases; ++basis)
        {
        //  if (evalStyle_ == panzer::EvaluatorStyle::EVALUATES)
          field_(cell, basis) = 0.0;
          for (int qp(0); qp < numQP; ++qp) {
            for (int dim(0); dim<numDims; ++dim ) {
              field_(cell, basis) += M_diff_(cell,qp)* muGrad_(cell, qp, dim) *
                weightedGradBasis(cell, basis, qp, dim);
            }
            field_(cell, basis) += cdot_(cell, qp)* weightedBasis(cell, basis, qp);
        //    std::cout << qp << "," << m_freeEnergy(cell, qp) << ", " << weightedBasis(cell, basis, qp) << std::endl;
          }
        //  std::cout << cell << "," << basis << ", " << field_(cell, basis) << std::endl;
        } // end loop over the basis functions
    }

} // end of evaluateFields()



#endif 
